package xin.marcher.wind.migrate.config;

import com.alibaba.druid.pool.DruidDataSource;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.shardingsphere.driver.api.ShardingSphereDataSourceFactory;
import org.apache.shardingsphere.infra.config.algorithm.AlgorithmConfiguration;
import org.apache.shardingsphere.sharding.api.config.ShardingRuleConfiguration;
import org.apache.shardingsphere.sharding.api.config.rule.ShardingTableRuleConfiguration;
import org.apache.shardingsphere.sharding.api.config.strategy.sharding.ComplexShardingStrategyConfiguration;
import org.apache.shardingsphere.sharding.api.sharding.complex.ComplexKeysShardingAlgorithm;
import org.springframework.boot.jdbc.DataSourceBuilder;
import xin.marcher.wind.migrate.sharding.DataSourceProperties;
import xin.marcher.wind.migrate.sharding.MigrateDataSourceProperties;
import xin.marcher.wind.migrate.sharding.TableRuleProperties;

import javax.sql.DataSource;
import java.sql.SQLException;
import java.util.*;

/**
 * 数据源配置
 */
public abstract class AbstractDataSourceConfig {

    /**
     * 根据配置构建数据源
     */
    public DataSource buildDataSource(MigrateDataSourceProperties migrateDataSourceProperties) throws SQLException {
        return CollectionUtils.isNotEmpty(migrateDataSourceProperties.getTableRules()) ?
                buildShardingDataSource(migrateDataSourceProperties) : buildDruidDataSource(migrateDataSourceProperties.getDataSources().get(0));
    }

    /**
     * 构建 Druid 数据源
     */
    public DruidDataSource buildDruidDataSource(DataSourceProperties dataSourceProperties) {
        DruidDataSource druidDataSource = DataSourceBuilder.create()
                .type(DruidDataSource.class)
                .driverClassName(dataSourceProperties.getDriverClassName())
                .url(dataSourceProperties.getUrl())
                .username(dataSourceProperties.getUsername())
                .password(dataSourceProperties.getPassword())
                .build();
        druidDataSource.setTestOnBorrow(true);
        druidDataSource.setTestWhileIdle(true);
        return druidDataSource;
    }

    /**
     * 构建 shardingsphere 数据源
     */
    private DataSource buildShardingDataSource(MigrateDataSourceProperties migrateDataSourceProperties) throws SQLException {
        // 1. 配置数据源
        Map<String, DataSource> dataSourceMap = buildDataSourceMap(migrateDataSourceProperties.getDataSources());
        // 2. 配置分片规则
        ShardingRuleConfiguration shardingRuleConfig = buildShardingRuleConfiguration(migrateDataSourceProperties);
        // 3. 配置其他的属性
        Properties properties = new Properties();
        properties.put("sql.show", migrateDataSourceProperties.getSqlshow());
        // 创建 ShardingSphereDataSource
        return ShardingSphereDataSourceFactory.createDataSource(dataSourceMap, Collections.singletonList(shardingRuleConfig), properties);
    }

    /**
     * 构建配置数据源
     */
    private Map<String, DataSource> buildDataSourceMap(List<DataSourceProperties> dataSourcePropertiesList) {
        Map<String, DataSource> dataSourceMap = new HashMap<>();
        for (int i = 0; i < dataSourcePropertiesList.size(); i++) {
            // 获取目标数据源配置
            DataSourceProperties dataSourceProperties = dataSourcePropertiesList.get(i);
            // 根据配置创建 dataSource
            DruidDataSource dataSource = buildDruidDataSource(dataSourceProperties);
            // 将目标数据源放入 dataSourceMap
            dataSourceMap.put("ds" + i, dataSource);
        }
        return dataSourceMap;
    }

    /**
     * 配置分片规则
     */
    private ShardingRuleConfiguration buildShardingRuleConfiguration(MigrateDataSourceProperties migrateDataSourceProperties) {
        // 分片规则
        ShardingRuleConfiguration shardingRuleConfig = new ShardingRuleConfiguration();

        // 最后一个 db 的下标
        int dbLastIndex = migrateDataSourceProperties.getDataSources().size() - 1;
        // 最后一个 table 的下标
        int tableLastIndex = migrateDataSourceProperties.getTableNum() - 1;

        for (TableRuleProperties tableRule : migrateDataSourceProperties.getTableRules()) {
            //构建分片策略实例
            ComplexKeysShardingAlgorithm<Comparable<?>> dbShardingAlgorithm = buildAlgorithmInstance(tableRule.getDbShardingAlgorithm());
            String dbShardingAlgorithmName = dbShardingAlgorithm.getClass().getName();
            ComplexKeysShardingAlgorithm<Comparable<?>> tableShardingAlgorithm = buildAlgorithmInstance(tableRule.getTableShardingAlgorithm());
            String tableShardingAlgorithmName = tableShardingAlgorithm.getClass().getName();

            // 配置表规则和分库分表策略
            ShardingTableRuleConfiguration tableRuleConfiguration = new ShardingTableRuleConfiguration(tableRule.getLogicTable(), "ds${0.." + dbLastIndex + "}." + tableRule.getLogicTable() + "${0.." + tableLastIndex + "}");
            tableRuleConfiguration.setDatabaseShardingStrategy(new ComplexShardingStrategyConfiguration(tableRule.getDbShardingColumns(), dbShardingAlgorithmName));
            tableRuleConfiguration.setTableShardingStrategy(new ComplexShardingStrategyConfiguration(tableRule.getTableShardingColumns(), tableShardingAlgorithmName));
            // 配置分片规则
            shardingRuleConfig.getTables().add(tableRuleConfiguration);

            // algorithmName 自行指定，需要和分片策略中的分片算法一致
            Properties dbProperties = new Properties();
            dbProperties.put("strategy", "COMPLEX");
            dbProperties.put("algorithmClassName", dbShardingAlgorithmName);
            shardingRuleConfig.getShardingAlgorithms().put(dbShardingAlgorithmName, new AlgorithmConfiguration("CLASS_BASED", dbProperties));
            Properties tableProperties = new Properties();
            tableProperties.put("strategy", "COMPLEX");
            tableProperties.put("algorithmClassName", tableShardingAlgorithmName);
            shardingRuleConfig.getShardingAlgorithms().put(tableShardingAlgorithmName, new AlgorithmConfiguration("CLASS_BASED", tableProperties));
        }
        return shardingRuleConfig;
    }

    /**
     * 构建分片策略实例
     */
    @SuppressWarnings("unchecked")
    private ComplexKeysShardingAlgorithm<Comparable<?>> buildAlgorithmInstance(String shardingAlgorithm) {
        try {
            Class<?> algorithmClass = Class.forName(shardingAlgorithm);
            return (ComplexKeysShardingAlgorithm<Comparable<?>>) algorithmClass.newInstance();
        } catch (Exception e) {
            throw new RuntimeException("build shardingAlgorithm instance error");
        }
    }

}
